package com.hliushi.netty.protocol;

import com.esotericsoftware.kryo.Kryo;
import com.esotericsoftware.kryo.io.Input;
import com.esotericsoftware.kryo.io.Output;
import com.esotericsoftware.kryo.serializers.BeanSerializer;
import org.apache.commons.io.IOUtils;

import java.io.ByteArrayInputStream;
import java.io.ByteArrayOutputStream;

/**
 * @author llhuang10
 * @date 2021/10/27 15:22
 */
public class KryoSerializer implements Serializer {

    private final Class<?> clazz;

    final ThreadLocal<Kryo> kryoThreadLocal = new ThreadLocal<Kryo>() {
        @Override
        protected Kryo initialValue() {
            Kryo kryo = new Kryo();
            kryo.register(clazz, new BeanSerializer<>(kryo, clazz));
            return kryo;
        }
    };

    public KryoSerializer(Class<?> aClass) {
        this.clazz = aClass;
    }

    private Kryo getKryo() {
        return kryoThreadLocal.get();
    }

    @Override
    public <T> T deserialize(Class<T> clazz, byte[] bytes) {
        if (bytes == null) {
            return null;
        }

        ByteArrayInputStream byteArrayInputStream = new ByteArrayInputStream(bytes);
        Input input = new Input(byteArrayInputStream);

        try {
            Kryo kryo = this.getKryo();
            return (T) kryo.readObjectOrNull(input, clazz);
        } finally {
            IOUtils.closeQuietly(input);
            IOUtils.closeQuietly(byteArrayInputStream);
        }

    }

    @Override
    public <T> byte[] serialize(T object) {
        ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
        Output output = new Output(byteArrayOutputStream);

        try {
            Kryo kryo = this.getKryo();
            kryo.writeObjectOrNull(output, object, object.getClass());
            output.flush();
            return byteArrayOutputStream.toByteArray();
        } finally {
            IOUtils.closeQuietly(output);
            IOUtils.closeQuietly(byteArrayOutputStream);
        }

    }
}
